package learning.pcfg.training;

import java.text.DecimalFormat;

import learning.data.Dataset;
import learning.data.document.TreeDocument;
import learning.pcfg.inference.IParseScorer;
import learning.pcfg.training.BasicParseScorer.Type;
import learning.pcfg.training.PCFG.DecodingResult;

public class CollinsTraining {

	public int maxIterations = 10;
	public float avgIterationsProportion = 0.75f;
	public boolean computeAvgParameters = true;
	
	IParseScorer allParseScorer, trueParseScorer;
	
	public CollinsTraining() {
		this.allParseScorer = new BasicParseScorer(Type.ALL_TERMINALS);
		this.trueParseScorer = new BasicParseScorer(Type.TRUE_TERMINALS);
	}
	
	public void test(Dataset<TreeDocument> testData, PCFGParameters parameters) {
		Evaluator evaluator = new Evaluator(testData.numLabels());

		for (TreeDocument doc : testData) {
			
			// compute most likely label under current parameters
			DecodingResult predictedParse = PCFG.decode(doc, allParseScorer, parameters);
			DecodingResult trueParse = PCFG.decode(doc, trueParseScorer, parameters);
			
			// update evaluator
			evaluator.update(predictedParse, trueParse);
		}
		
		evaluator.print();		
	}
	
	public PCFGParameters train(Dataset<TreeDocument> trainingData) {
		
		PCFGParameters avgParameters = new PCFGParameters();
		PCFGParameters iterAvgParameters = new PCFGParameters();
		PCFGParameters iterParameters = new PCFGParameters();
		
		int avgIterations = maxIterations - (int)(maxIterations*avgIterationsProportion);
		
		Evaluator evaluator = new Evaluator(trainingData.numLabels());
		
		for (int i = 0; i < maxIterations; i++)
		{
			evaluator.reset();
			
			float delta = (float)(2.0 / (2.0 + i));
			boolean useIterAverage = computeAvgParameters && (i >= maxIterations - avgIterations);
			
			trainingIteration(trainingData, delta, evaluator, iterParameters, useIterAverage, iterAvgParameters);

			evaluator.print();

			if (computeAvgParameters && useIterAverage) {
				avgParameters.sum(iterAvgParameters, 1.0f / avgIterations);
			}
		}
		evaluator.print();
		
		return (computeAvgParameters)? avgParameters : iterParameters;
	}
	
	public void trainingIteration(Dataset<TreeDocument> trainingData, float delta, Evaluator evaluator,
			PCFGParameters iterParameters, boolean useIterAverage, PCFGParameters iterAvgParameters) {
		
		// resetting parameters to 0
		iterAvgParameters.reset();
		iterParameters.reset();
		
		trainingData.randomizeOrder();
		
		for (TreeDocument doc : trainingData) {
			
			// compute most likely labels under current parameters
			
			DecodingResult predictedParse = PCFG.decode(doc, allParseScorer, iterParameters);
			
			// compute most likely parse for the correct labels under current parameters
			
			DecodingResult trueParse = PCFG.decode(doc, trueParseScorer, iterParameters);
			
			// update weights
			PCFG.update(doc, predictedParse, trueParse, iterParameters, delta);

			if (useIterAverage) {				
				iterAvgParameters.sum(iterParameters, 1.0f / trainingData.numDocs());
			}
			
			// update evaluator
			evaluator.update(predictedParse, trueParse);
		}
	}
	
	public static class Evaluator {
		
		int numLabels;
		int[][] confusionMatrix;
		DecimalFormat format;
		String ansiRed    = "\u0027[31m";
		String ansiNormal = "\u0027[0m";
		boolean useAnsi = true;
		
		public Evaluator(int numLabels) {
			this.numLabels = numLabels;
			this.confusionMatrix = new int[numLabels][numLabels];
			this.format = new DecimalFormat("0.000");
		}

		public void reset() {
			for (int i=0; i < numLabels; i++)
				for (int j=0; j < numLabels; j++)
					confusionMatrix[i][j] = 0;
		}		
		
		public void update(DecodingResult predictedParse, DecodingResult trueParse) {
			for (int i=0; i < predictedParse.labels.length; i++)
				confusionMatrix[trueParse.labels[i]][predictedParse.labels[i]]++;
		}
		
		public void print() {
			int[] trueCounts = new int[numLabels]; 
			for (int i=0; i < numLabels; i++)
				for (int j=0; j < numLabels; j++)
					trueCounts[i] += confusionMatrix[i][j];
			
			StringBuilder sb = new StringBuilder();
			for (int i=0; i < numLabels; i++) {
				sb.append("\t" + i);
			}
			sb.append("\n");
			
			for (int i=0; i < numLabels; i++) {
				sb.append(i);
				
				for (int j=0; j < numLabels; j++) {
					sb.append("\t");
					float f = (float)(confusionMatrix[i][j] / (double)trueCounts[i]);
					String n = format.format(f);
					if (f > 0 && useAnsi) n = ansiRed + n + ansiNormal;
					sb.append(n);
				}
				sb.append("\n");
			}
			System.out.println(sb.toString());
		}
	}
}
